from typing import Optional

from python.mypackage.TreeNode import TreeNode


class Solution:
    """
    方法：countNodes 计算完全二叉树的节点数量

    Args:
        root: Optional[TreeNode] - 完全二叉树的根节点

    Returns:
        int: 完全二叉树的节点总数

    Time: O(logN * logN) - 其中N是节点数，每次递归调用get_depth需要O(logN)时间，递归深度为O(logN)

    Space: O(logN) - 递归调用栈的深度
    """
    def countNodes(self, root: Optional[TreeNode]) -> int:
         if not root: return 0
         left_depth = self.get_depth(root.left)
         right_depth = self.get_depth(root.right)
         if left_depth == right_depth:
            # 左子树是满二叉树，可以直接计算其节点数(2^left_depth)
            # 然后递归计算右子树
            return (1 << left_depth) + self.countNodes(root.right)
         else:
            # 右子树是满二叉树(比左子树少一层)
            # 递归计算左子树
            return (1 << right_depth) + self.countNodes(root.left)

    def get_depth(self, node: Optional[TreeNode]) -> int:
        depth = 0
        while node:
            depth += 1
            node = node.left
        return depth

    
    def countNodes1(self, root: Optional[TreeNode]) -> int:
        if not root: return 0
        node = root
        level = 0
        while node.left:
            level += 1
            node = node.left
        low = 1 << level
        high = (1 << (level + 1)) - 1
        while low < high:
            mid = (high - low + 1) // 2 + low
            if self.exists(root, level, mid):
                low = mid
            else:
                high = mid - 1
        return low

    def exists(self, root: Optional[TreeNode], level: int, k: int) -> bool:
        bits = 1 << (level - 1)
        node = root
        while node and bits > 0:
            if (bits & k) == 0:
                node = node.left
            else:
                node = node.right
            bits >>= 1
        return node is not None